package org.easy.byteaidraw;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import org.easy.core.exception.BizRuntimeException;
import org.easy.core.util.FunctionThreadQueueWorker;
import org.easy.byteaidraw.dto.AnimeDrawRequest;
import org.easy.byteaidraw.dto.ByteAiDrawErrModel;
import org.easy.byteaidraw.dto.RoleFeatureRetentionRequest;
import com.volcengine.service.visual.IVisualService;
import com.volcengine.service.visual.impl.VisualServiceImpl;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.List;


/**
 * 通用方法封装工具类
 */
@Slf4j
public class ByteAiDrawClient {


    public ByteAiDrawClient(String accessKey, String secretKey,
                            Integer qps1,
                            Integer qps2,
                            Integer qps3
    ) {
        this.accessKey = accessKey;
        this.secretKey = secretKey;
        this.functionThreadQueueWorker = new FunctionThreadQueueWorker(qps1);
        this.functionThreadQueueWorker2 = new FunctionThreadQueueWorker(qps2);
        this.functionThreadQueueWorker3 = new FunctionThreadQueueWorker(qps3);
    }

    private final String accessKey;
    private final String secretKey;
    private final FunctionThreadQueueWorker functionThreadQueueWorker;
    private final FunctionThreadQueueWorker functionThreadQueueWorker2;
    private final FunctionThreadQueueWorker functionThreadQueueWorker3;


    /**
     * https://www.volcengine.com/docs/6791/1213131
     *
     * @param request
     * @return
     */
    public String animeDrawSubmit(AnimeDrawRequest request) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);


        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "high_aes");
        req.put("prompt", request.getPrompt());
        req.put("model_version", "anime_v1.3.1");
        if (StringUtils.isNotEmpty(request.getOriginImg())) {
            req.put("image_urls", Lists.newArrayList(request.getOriginImg()));
        }
        //0.0<strength=<1.0之间，（越小越接近原图，越大越接近文本控制)
        req.put("strength", request.getStrength());
        req.put("scale", request.getScale());
        req.put("ddim_steps", 50);
        req.put("width", request.getWidth());
        req.put("height", request.getHeight());


        try {
            long start = System.currentTimeMillis();
            Object enqueue = functionThreadQueueWorker.enqueue(o -> {
                try {
                    return visualService.cvSync2AsyncSubmitTask(req);
                } catch (Exception e) {
                    return new ByteAiDrawErrModel(50412, e.getMessage());
                }
            }, req);

            log.info("animeDraw消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));

            JSONObject response = (JSONObject) enqueue;
            if (response.getInteger("code") == 10000) {
                return response.getJSONObject("data").getString("task_id");
            }
            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }


    public String animeDrawGetResult(String taskId) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);


        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "high_aes");
        req.put("task_id", taskId);
        req.put("req_json", "{\"return_url\":true}");

        try {
            long start = System.currentTimeMillis();
            JSONObject response = (JSONObject) visualService.cvSync2AsyncGetResult(req);
            log.info("animeDrawGetResult消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));
            if (response.getInteger("code") == 10000) {
                JSONArray jsonArray = response.getJSONObject("data").getJSONArray("image_urls");
                if (jsonArray == null) {
                    return null;
                }
                return jsonArray.getString(0);
            }
            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }


    /**
     * 通用2.0Pro-指令编辑 https://www.volcengine.com/docs/6791/1384311
     * @param prompt
     * @param originImg  编辑指令使用自然语言即可
     * 每次编辑使用单指令会更好
     * 局部编辑时指令描述尽量精准，尤其是画面有多个实体的时候，描述清楚对谁做什么，能获取更精准的编辑效果
     * 发现编辑效果不明显的时候，可以调整一下编辑强度scale，数值越大越贴近指令执行
     * 尽量使用清晰的，分辨率高的底图，动漫及豆包模型生成的图片编辑效果会更好
     *
     *                   参考示例
     * 添加/删除实体：添加/删除xxx（删除图上的女孩/添加一道彩虹）
     * 修改实体：把xxx改成xxx（把手里的鸡腿变成汉堡）
     * 修改风格：改成xxx风格（改成漫画风格）
     * 修改色彩：把xxx改成xx颜色（把衣服改成粉色的）
     * 修改动作：修改表情动作（让他哭/笑/生气）
     * 修改环境背景：背景换成xxx，在xxx（背景换成海边/在星空下）
     *
     * @param scale 影响文本描述的程度，该值越大代表文本描述影响程度越大，且输入图片影响程度越小 默认值：0.5 取值范围：[0, 1]
     * @return
     */
    public String byteedit_v2_0(String prompt, String originImg, Float scale) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);

        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "byteedit_v2.0");
        req.put("image_urls", Lists.newArrayList(originImg));
        req.put("prompt", prompt);

        req.put("scale", scale == null ? 0.5f : scale);
        req.put("return_url", true);
        req.put("use_sr", true);

        try {
            long start = System.currentTimeMillis();
            Object enqueue = functionThreadQueueWorker2.enqueue(o -> {
                try {
                    return visualService.cvSync2AsyncSubmitTask(req);
                } catch (Exception e) {
                    return new ByteAiDrawErrModel(50412, e.getMessage());
                }
            }, req);

            log.info("animeDraw消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));

            JSONObject response = (JSONObject) enqueue;
            if (response.getInteger("code") == 10000) {
                return response.getJSONObject("data").getString("task_id");
            }

            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }

    /**
     * 角色特征保持 https://www.volcengine.com/docs/6791/1361423
     *
     * @param request
     */
    public String roleFeatureRetention(RoleFeatureRetentionRequest request) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);

        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "high_aes_ip_v20");
        req.put("image_urls", Lists.newArrayList(request.getOriginImg()));
        req.put("prompt", request.getPrompt());
        req.put("desc_pushback", true);
        req.put("scale", request.getScale());
        req.put("ddim_steps", request.getDdimSteps());
        req.put("width", request.getWidth());
        req.put("height", request.getHeight());
        req.put("cfg_rescale", 0.7);
        req.put("ref_ip_weight", request.getRef_ip_weight());
        req.put("ref_id_weight", request.getRef_id_weight());
        req.put("return_url", true);
        req.put("use_sr", true);

        try {
            long start = System.currentTimeMillis();
            Object enqueue = functionThreadQueueWorker2.enqueue(o -> {
                try {
                    return visualService.cvSync2AsyncSubmitTask(req);
                } catch (Exception e) {
                    return new ByteAiDrawErrModel(50412, e.getMessage());
                }
            }, req);

            log.info("animeDraw消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));

            JSONObject response = (JSONObject) enqueue;
            if (response.getInteger("code") == 10000) {
                return response.getJSONObject("data").getString("task_id");
            }

            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }

    public String waitForResultByRoleFeatureRetention(String generateUuid) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);


        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "high_aes_ip_v20");
        req.put("task_id", generateUuid);
        req.put("req_json", "{\"return_url\":true}");

        try {
            long start = System.currentTimeMillis();
            JSONObject response = (JSONObject) visualService.cvSync2AsyncGetResult(req);
            log.info("waitForResultByRoleFeatureRetention消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));
            if (response.getInteger("code") == 10000) {
                JSONArray jsonArray = response.getJSONObject("data").getJSONArray("image_urls");
                if (jsonArray == null) {
                    return null;
                }
                return jsonArray.getString(0);
            }
            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }


    /**
     * https://www.volcengine.com/docs/6793/1324393#%E8%AF%B7%E6%B1%82-%E8%BF%94%E5%9B%9E%E5%AE%8C%E6%95%B4%E7%A4%BA%E4%BE%8B
     *
     * @param prompt
     * @param originImgs
     * @param width
     * @param height
     * @return
     */
    public String img2img_anime_accelerated_maintain_id_for_smart_drawing_anime(String prompt, List<String> originImgs, Long width, Long height) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);

        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "img2img_anime_accelerated_maintain_id_for_smart_drawing_anime");

        req.put("positive_prompt", "best_quality,anime_style," + prompt);

        req.put("face_image", "uri://binary_data?index=0");
        req.put("style_image", "uri://binary_data?index=1");
        req.put("image_urls", originImgs);


        req.put("face_tagger_switch", true);
        req.put("face_v2_switch", true);
        req.put("style_switch", false);

        req.put("face_detection_switch", true);
        req.put("step", 50);
        req.put("cfg", 6.0f);
        req.put("facestyle_switch", false);
        req.put("width", width);
        req.put("height", height);
        req.put("clarity", 0.6);
        req.put("return_url", true);
//        req.put("pose_switch", true);

        try {
            long start = System.currentTimeMillis();
            Object enqueue = functionThreadQueueWorker3.enqueue(o -> {
                try {
                    return visualService.cvSync2AsyncSubmitTask(req);
                } catch (Exception e) {
                    return new ByteAiDrawErrModel(50412, e.getMessage());
                }
            }, req);

            log.info("animeDraw消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));

            JSONObject response = JSON.parseObject(JSON.toJSONString(enqueue));
            if (response.getInteger("code") == 10000) {
                return response.getJSONObject("data").getString("task_id");
            }
            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }

    public String waitForResultByimg2img_anime_accelerated_maintain_id_for_smart_drawing_anime(String generateUuid) {
        IVisualService visualService = VisualServiceImpl.getInstance();
        visualService.setAccessKey(accessKey);
        visualService.setSecretKey(secretKey);


        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "img2img_anime_accelerated_maintain_id_for_smart_drawing_anime");
        req.put("task_id", generateUuid);
        req.put("req_json", "{\"return_url\":true}");

        try {
            long start = System.currentTimeMillis();
            JSONObject response = (JSONObject) visualService.cvSync2AsyncGetResult(req);
            log.info("waitForResultByRoleFeatureRetention消耗时间{} 请求参数{}", (System.currentTimeMillis() - start), JSON.toJSONString(req));
            if (response.getInteger("code") == 10000) {
                JSONArray jsonArray = response.getJSONObject("data").getJSONArray("image_urls");
                if (jsonArray == null) {
                    return null;
                }
                return jsonArray.getString(0);
            }
            throw new BizRuntimeException(response.getString("message"));
        } catch (Exception e) {
            throw new BizRuntimeException(e.getMessage());
        }
    }


}